import os
from dotenv import load_dotenv
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_community.llms import Tongyi
from langchain.embeddings.dashscope import DashScopeEmbeddings

# 加载环境变量
load_dotenv()

def main():
    # 确保设置了通义千问 API密钥
    if not os.getenv("DASHSCOPE_API_KEY"):
        print("请设置DASHSCOPE_API_KEY环境变量或在.env文件中配置")
        return
    
    # 示例文本
    documents = [
        "LangChain是一个用于开发由语言模型驱动的应用程序的框架。它可以帮助应用程序更具有上下文感知能力。",
        "LangChain提供了许多模块，可用于各种应用程序。这些模块可以组合起来创建更复杂的应用程序，或者单独使用。",
        "LangChain的主要组件包括链（Chains）、代理（Agents）、记忆（Memory）、提示（Prompts）等。",
        "LangChain支持与多种语言模型集成，包括OpenAI、Anthropic、Hugging Face等提供的模型。"
    ]
    
    # 文本分割
    text_splitter = CharacterTextSplitter(chunk_size=200, chunk_overlap=0)
    texts = text_splitter.create_documents(documents)
    
    # 创建向量存储
    embeddings = DashScopeEmbeddings(model="text-embedding-v1")
    vectorstore = Chroma.from_documents(texts, embeddings)
    
    # 创建检索QA链
    llm = Tongyi(temperature=0, model_name="qwen-turbo")
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vectorstore.as_retriever()
    )
    
    # 进行问答
    print("\n欢迎使用LangChain文档问答系统! 输入'退出'结束对话。\n")
    print("您可以询问关于LangChain的问题，系统会基于提供的文档回答。\n")
    
    while True:
        query = input("\n问题: ")
        if query.lower() in ["退出", "exit", "quit"]:
            print("\n谢谢使用，再见!")
            break
        
        # 获取回复
        result = qa_chain.invoke({"query": query})
        print(f"\n回答: {result['result']}")

if __name__ == "__main__":
    main()